# -*- coding: utf-8 -*-
"""
ToxoFundus 数据集预处理脚本
---------------------------------
功能：
    - 处理弓形虫性脉络膜视网膜炎 (ToxoFundus) 六分类数据集；
    - 每类图像存放于独立文件夹中；
    - 输出统一格式的 annotations.json 与统计文件。

使用样例：
    python src/toxofundus_preprocess.py \
        --data_path /home/zhangpinglu/data0/gy/Dataset/public_dataset/ToxoFundus \
        --tar_path /home/zhangpinglu/data0/gy/Dataset/public_processed/ToxoFundus \
        --prefix toxofundus
"""

import os
import sys
import json
import argparse
from tqdm import tqdm
from PIL import Image
from data_preprocess.utils.crop import crop_resize_save
from data_preprocess.utils.stat import save_classification_stats

# 六分类映射表
TOXOFUNDUS_LABEL_MAP = {
    "healthy": "Normal fundus appearance (no retinal lesions detected)",
    "active": "Active toxoplasmosis lesion (ongoing inflammation and retinal exudation)",
    "inactive": "Inactive toxoplasmosis scar (healed lesion with chorioretinal atrophy)",
    "active_inactive": "Mixed active and inactive toxoplasmosis lesions (coexistence of old scar and active inflammation)",
    "active_active": "Multiple active toxoplasmosis lesions (multifocal active chorioretinitis)",
    "inactive_inactive": "Multiple inactive toxoplasmosis scars (multifocal healed lesions without active inflammation)"
}


def gather_data(data_path, tar_path, prefix="toxofundus", resize=(512, 512)):
    """
    对 ToxoFundus 数据集进行统一预处理。
    步骤：
        1. 遍历六个类别；
        2. 对每张图片执行裁剪与 resize；
        3. 构建 annotations.json；
        4. 生成统计文件。

    参数：
        data_path: 原始数据根路径（包含 Data_Raw_6class_All/Data_Raw_6class_All）
        tar_path: 输出路径
        prefix: 输出图片名前缀
        resize: 目标尺寸 (宽, 高)
    """
    raw_root = os.path.join(data_path, "Data_Raw_6class_All", "Data_Raw_6class_All")
    if not os.path.exists(raw_root):
        raise FileNotFoundError(f"❌ 数据路径不存在: {raw_root}")

    os.makedirs(tar_path, exist_ok=True)
    image_dir = os.path.join(tar_path, "images")
    os.makedirs(image_dir, exist_ok=True)
    os.makedirs("./experiments/stat", exist_ok=True)

    data_dict = {}
    image_id = 0

    for cls_folder, cls_desc in TOXOFUNDUS_LABEL_MAP.items():
        cls_dir = os.path.join(raw_root, cls_folder)
        if not os.path.exists(cls_dir):
            print(f"⚠️ 跳过 {cls_folder}: 路径不存在")
            continue
        image_list = sorted(os.listdir(cls_dir))
        for image_name in tqdm(image_list, desc=f"处理 {cls_folder}", unit="img"):
            if not image_name.lower().endswith((".jpg", ".jpeg", ".png", ".tif")):
                continue

            src_path = os.path.join(cls_dir, image_name)
            image_id += 1
            new_name = f"{prefix}_{image_id:06d}.png"
            dst_path = os.path.join(image_dir, new_name)

            try:
                crop_info = crop_resize_save(
                    image_path=src_path,
                    save_path=dst_path,
                    resize=resize,
                    crop_threshold=25
                )
            except Exception as e:
                raise RuntimeError(f"❌ 图像处理失败 {src_path}: {e}")

            entry = {
                "image_name": new_name,
                "image_path": dst_path,
                "original_path": src_path,
                "crop_info": crop_info,
                "diagnosis": {
                    "classification": {"text": cls_desc}
                }
            }
            data_dict[new_name] = entry

    # 保存 annotations.json
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(data_dict, f, indent=4)

    # 保存统计
    save_classification_stats(data_dict, "./experiments/stat/toxofundus.json")

    print(f"✅ 预处理完成，共处理 {len(data_dict)} 张图片。输出位于 {tar_path}")
    return data_dict


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ToxoFundus 六分类数据集预处理")
    parser.add_argument("--data_path", type=str,
                        default="/home/zhangpinglu/data0/gy/Dataset/public_dataset/ToxoFundus",
                        help="原始 ToxoFundus 数据集根目录")
    parser.add_argument("--tar_path", type=str,
                        default="/home/zhangpinglu/data0/gy/Dataset/public_processed/ToxoFundus",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="toxofundus",
                        help="输出图片名称前缀")
    parser.add_argument("--resize", type=int, nargs=2, default=[512, 512],
                        help="目标尺寸 (宽, 高)")
    args = parser.parse_args()

    gather_data(args.data_path, args.tar_path, prefix=args.prefix, resize=tuple(args.resize))
